Skip to content

feat: add ASTTransformer for in-place AST node replacement#80

Merged
boolangery merged 1 commit into
masterfrom
feat/ast-transformer
Jul 20, 2026
Merged

feat: add ASTTransformer for in-place AST node replacement#80
boolangery merged 1 commit into
masterfrom
feat/ast-transformer

Conversation

@boolangery

Copy link
Copy Markdown
Owner

Problem

The existing ASTVisitor and ASTRecursiveVisitor only support read-only traversal. There is no way to transform or replace AST nodes during traversal, which is needed for tasks like:

  • Constant folding / expression simplification
  • Variable renaming
  • Code deobfuscation
  • AST-based optimization passes

Solution

Adds ASTTransformer, a production-grade AST visitor that supports in-place node replacement.

Design

  • Iterative stack-based traversal — avoids recursion limits on deep trees
  • Direct parent tracking via (key, container) tuples — O(1) replacement, no brute-force key search
  • Replacement children visited — after replacing a node, children of the new node are traversed
  • Identity checkingreplacement is not node avoids unnecessary mutation
  • Root replacement — the transformed root is returned from visit()

Usage

class NumberDoubler(ASTTransformer):
    def visit_Number(self, node):
        return Number(node.n * 2)

tree = ast.parse("x = 5")
new_tree = NumberDoubler().visit(tree)

Comparison with PR #74

This implementation improves on the original PR by:

  • O(1) parent replacement via tracked keys (vs O(n×m) brute-force dict search)
  • Proper root replacement support
  • Visiting replacement node's children (not stale original children)
  • Comprehensive test suite (15 tests)
  • No commented-out debug code
  • Clean docstrings

Tests

15 tests covering:

  • Value replacement, statement replacement, deep nesting
  • Root replacement, recursive constant folding
  • Round-trip identity, None/no-op semantics
  • List child iteration, parent class dispatch

Closes #74

Adds ASTTransformer, an iterative depth-first visitor that supports
replacing AST nodes during traversal. Unlike the existing ASTVisitor,
the transformer propagates replacements back to parent nodes and
supports root replacement.

Key design decisions:
- Stack-based iterative traversal (no recursion limit issues)
- Direct parent tracking via (key, container) tuples (O(1) replacement)
- Replacement node's children are visited (not original's)
- Identity check (replacement is not node) avoids unnecessary writes
- Returns new root from visit(), enabling full tree transformation

Includes 15 comprehensive tests covering:
- Value replacement, statement replacement, deep nesting
- Root replacement, recursive folding, round-trip identity
- None/no-op semantics, list child iteration, parent class dispatch
@boolangery
boolangery merged commit 797d41a into master Jul 20, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant